from sklearn.mixture import GaussianMixture
import numpy as np
import cuml
from tqdm import tqdm
# 用BIC搜参数,只是确定参数的话，可以考虑每次随机采样一点，就五千吧。
import time
from ..config import chunk_size,context_size







def bic(X):
    X=X.numpy()
    # 数量太少不需要聚类了。
    if X.shape[0]<=1:
        return [0 for _ in range(X.shape[0])]
    
    # 主要考虑到如果小于这个分块数量，那一个chunk内的内容会超过上下文长度。
    cluster_upper_bound = min(130,X.shape[0])
    cluster_lower_bound= max(2,cluster_upper_bound-10)
    print('cluster_lower_bound',cluster_lower_bound)
    
    n_components_range = range(cluster_lower_bound,cluster_upper_bound,200)
    # n_components_range = range(125,126,10)
    n_neighbors_upper_bound = 91
    n_neighbors_range = range(90,n_neighbors_upper_bound,30) 

    # 存储每个参数的BIC值
    bic_values = [[float('inf')  for _ in range(n_neighbors_upper_bound)] for _ in range(cluster_upper_bound)]
    
    labels = [[float('inf')  for _ in range(n_neighbors_upper_bound)] for _ in range(cluster_upper_bound)]


    # 遍历每个组件数量
    for n_components in n_components_range:
        for n_neighbors in n_neighbors_range:
            
            start=time.time()
            try:
                clusterable_embedding = cuml.UMAP(
                    n_neighbors=n_neighbors, # 如果超过X的数量会被自动截断到X-1
                    min_dist=0.0,
                    n_components=16,
                    # random_state=42,
                    # n_jobs=-1,
                    # low_memory=False,
                    # verbose=True,
                ).fit_transform(X)
            except Exception as e:
                print(f'UMAP降维失败了，当前的X形状是{X.shape}')
                # print(f'异常是：{e}')
                clusterable_embedding = X
            end=time.time()
            print(f'UMAP time cose{end-start}')
            # 初始化GMM模型
            gmm = GaussianMixture(n_components=n_components,init_params='k-means++')
            # 拟合模型并预测
            
            
            labels[n_components][n_neighbors] = gmm.fit_predict(clusterable_embedding)
            endd=time.time()
            print(f'GMM time cose{endd-end}')
            # 计算BIC
            bic_values[n_components][n_neighbors]=gmm.bic(clusterable_embedding)
           

    # 选择BIC最小的模型
    optimal_index = np.argmin(bic_values)  
    optimal_params = np.unravel_index(optimal_index, np.array(bic_values).shape)
    optimal_labels=labels[optimal_params[0]][optimal_params[1]]
    # print(bic_values)
    print(f"最优的参数是: {optimal_params}")
    # TODO 不要返回参数，直接返回最优的预测结果。
    return optimal_labels

if __name__=='__main__':
    # 测试
    from sklearn.datasets import make_blobs
    
    
    X = np.random.random((2,768))
    print(bic(X))
        